infile = r'mature.fa'
pre_file = r'remain_pre.fa'
outfile = r'discard_miRNA_from_miRbase2'
RNAfold_file = r'RNAfold_out'
from Bio import SeqIO
import numpy as np
import re
class miR(object):
    def __init__(self,name,pre):
        self.name = name
        self.pre= pre
        self.mat_5p = ""
        self.mat_3p = ""
        self.T=True
        self.reason = ""
        self.fold = ""
    def include(self,seq):
        if self.pre.find(seq)!=-1:
            start = self.pre.find(seq)
            if start < len(self.pre)/2:
                self.mat_5p = seq
            else:
                self.mat_3p = seq
    def check_5p3p(self):
        if self.mat_3p!="":
            l = len(self.mat_3p)
            if l<=18 or l>=26:
                self.T=False
                self.reason = 'miR-3p length out of 18~26 nt'
            return 0
        if self.mat_5p!="":
            l = len(self.mat_5p)
            if l<=18 or l>=26:
                self.T=False
                self.reason = 'miR-5p length out of 18~26 nt'
            return 0
        self.T = False
        self.reason = 'no miR-3p or miR-5p'
    def read_fold(self,fold_raw):
        k = re.compile('([(|)|.]+) [(](.+)[)]')
        k = k.findall(fold_raw)
        
        if len(k[0])>=2:

            self.fold = k[0][0]
            self.score = float(k[0][1])

    def read_RNAfold(self):
        # find two folder
        # start from miR5p
        if self.fold!='':
            # make a hill 
            self.hill = np.zeros(len(self.fold))
            height = 0
            for i in enumerate(self.fold):
                if i[1] == '(':
                    height+=1
                elif i[1] == ')':
                    height-=1
                self.hill[i[0]] = height

            d_3p = False
            if self.mat_5p!='':
                self.start_5p = self.pre.find(self.mat_5p)
                self.end_5p = self.start_5p + len(self.mat_5p)
                count_fold_5p = self.fold[self.start_5p:self.end_5p].count('(')
                if count_fold_5p<=16:
                    self.T=False
                    self.reason = 'miR-5p arm less than  complementarity '
                    return 0
                # digitalize the miR3p
                if self.mat_3p == "":
                    d_3p = True
                    # find end 
                    counter = 0
                    point = self.start_5p
                    while self.fold[point] != '(':
                        point+=1
                        counter +=1
                    height = self.hill[point]
                    self.end_3p = len(self.fold)
                    for i in range(self.end_5p+2,len(self.fold)):
                        if self.hill[i] < height:
                            self.end_3p = i+1+counter+2
                            break
                    # find start 
                    counter = 0
                    point = self.end_5p
                    while self.fold[point] != '(':
                        point-=1
                        counter -=1
                    height = self.hill[point]
                    for i in range(self.end_5p,len(self.fold)):
                        if self.hill[i] < height:
                            self.start_3p = i+1+counter+2
                            break
                    self.mat_3p = self.pre[self.start_3p:self.end_3p]


            if self.mat_3p!='':

                self.start_3p = self.pre.find(self.mat_3p)
                self.end_3p = self.start_3p + len(self.mat_3p)
                count_fold_3p = self.fold[self.start_3p:self.end_3p].count(')')
                if count_fold_3p<16 & d_3p == True:
                    self.T=False
                    self.reason = 'miR-3p arm less than  complementarity '
                    return 0
                if self.mat_5p == "":
                    # find end 
                    counter = 0
                    point = self.start_3p
                    while self.fold[point] != ')':
                        point+=1
                        counter +=1
                    height = self.hill[point]
                    for i in range(int(len(self.fold)/2),0,-1):
                        if self.hill[i] < height:
                            self.end_5p = i+1+counter+2
                            break
                    # find start 
                    counter = 0
                    point = self.end_5p
                    while self.fold[point] != ')':
                        point-=1
                        counter -=1
                    height = self.hill[point]
                    self.start_5p = 0
                    for i in range(int(len(self.fold)/2),0,-1):
                        if self.hill[i] < height:
                            self.start_5p = i+1+counter+2
                            break
                    self.mat_5p = self.pre[self.start_5p:self.end_5p]
            # the loop should large than 8
            if self.mat_5p!= '' and self.mat_3p != '':
                self.mat_3p_start = self.pre.find(self.mat_3p)
                self.mat_3p_end = self.pre.find(self.mat_3p)+len(self.mat_3p)
                self.mat_5p_start = self.pre.find(self.mat_5p)
                self.mat_5p_end = self.pre.find(self.mat_5p)+len(self.mat_5p)
                if self.mat_3p_start-self.mat_5p_end <8:
                    self.T=False
                    self.reason = "loop size less than 8nt"




                    # print '\n'
                    # print '-----\n'
                    # print self.fold
                    # print self.name,start_3p,end_3p
                    # print mat_3p
                    # print self.hill


                
dic = {}
for i in SeqIO.parse(pre_file,'fasta'):
    seq = str(i.seq).upper().replace('U','T')
    name = i.name
    dic[name.upper()] = miR(name,seq)

for i in SeqIO.parse(infile,'fasta'):
    k = str(i.name).split('-')
    seq = str(i.seq).upper().replace('U','T')
    if k[-1] in ['3p','5p','3P','5P','3p*','5p*','3P*','5P*']:
        name = '-'.join(k[:-1])
    else:
        name = str(i.name)
    if name.upper() in dic:
        dic[name.upper()].include(seq)
    else:
        pass

# with open(outfile,'w') as fila:
#     for i in dic:
#         dic[i].check_5p3p()
#         if dic[i].T!=True:
#             fila.write(dic[i].name+'\t'+dic[i].reason+'\n')

class RNAfold(object):
    def __init__(self,name):
        self.name = name
        self.seq = ''
        self.fold = ''
name = ""
seq = ""
fold = ""

t = 0
lista =[]
with open(RNAfold_file) as fila:
    for i in fila:
        i = i.replace('\n','')
        if len(i)>1:
            if i[0] == '>':
                name = i[1:].split(' ')[0]
                t = 1
                continue
            if t ==1:
                seq = i.upper().replace('U','T')
                t = 2
                continue
            if t ==2:
                fold = i
                k = RNAfold(name)
                k.seq = seq
                k.fold = fold
                lista.append(k)
                name = ''
                seq = ''
                fold = ''
if name!= '':
    k = RNAfold(name)
    k.seq = seq
    k.fold = fold
    lista.append(k)
 
for i in lista:
    if i.name.upper() in dic:
        dic[i.name.upper()].read_fold(i.fold)
        dic[i.name.upper()].read_RNAfold()
        dic[i.name.upper()].check_5p3p()


with open(outfile,'w') as fila:
    for i in dic:
        if dic[i].T ==False:
            fila.write(dic[i].name+'\t'+dic[i].reason+'\n')


